#!/usr/bin/env python3

import sys
sys.path.append('../../../tools')
from gmcore_plot import *

parser = argparse.ArgumentParser(description='Plot test result')
parser.add_argument('-i', '--input', help='Input data', default='cm.720x360x5.h0.nc')
parser.add_argument('-t', '--time-step', type=int, help='Time step to plot', default=-1)
parser.add_argument('-o', '--output', help='Output figure')
parser.add_argument('-l', '--loop', action='store_true', help='Loop plot')
args = parser.parse_args()

fig = plt.figure(figsize=(8, 8))
proj = ccrs.PlateCarree(central_longitude=180)

levels = np.linspace(-7e-5, 7e-5, 15)
colors = ListedColormap([
	[0.15, 0.24, 0.47, 1.00],
	[0.24, 0.40, 0.67, 1.00],
	[0.40, 0.61, 0.81, 1.00],
	[0.59, 0.78, 0.92, 1.00],
	[0.76, 0.89, 0.96, 1.00],
	[0.90, 0.95, 0.98, 1.00],
	[1.00, 1.00, 1.00, 0.00],
	[1.00, 1.00, 1.00, 0.00],
	[0.98, 0.95, 0.76, 1.00],
	[0.96, 0.82, 0.48, 1.00],
	[0.94, 0.64, 0.29, 1.00],
	[0.89, 0.44, 0.24, 1.00],
	[0.79, 0.25, 0.20, 1.00],
	[0.62, 0.16, 0.15, 1.00]
])

def plot(data_path, fig):
	f = addfile(data_path)

	if args.loop:
		ax = fig.add_subplot(2, 1, 1, projection=proj)
		ax.coastlines()
		days = ((f.time[args.time_step] - f.time[0]).values).astype("timedelta64[D]")
		plot_contour_map(ax, f.vor[args.time_step,0,:,:], cmap=colors, with_grid=True, with_contour=False, left_string=f'Vorticity (1/s) at {days}', levels=levels, add_cyclic_point=False)
	else:
		ax = fig.add_subplot(2, 1, 1, projection=proj)
		ax.coastlines()
		if f.time.size > 20:
			plot_contour_map(ax, f.vor[20,0,:,:], cmap=colors, with_grid=True, with_contour=False, left_string='Vorticity (1/s)', levels=levels, add_cyclic_point=False)
		if f.time.size > 40:
			plot_contour_map(ax, f.vor[40,0,:,:], cmap=colors, with_grid=True, with_contour=False, left_string='Vorticity (1/s)', levels=levels, add_cyclic_point=False)
		if f.time.size > 60:
			plot_contour_map(ax, f.vor[60,0,:,:], cmap=colors, with_grid=True, with_contour=False, left_string='Vorticity (1/s)', levels=levels, add_cyclic_point=False)
		if f.time.size > 80:
			plot_contour_map(ax, f.vor[80,0,:,:], cmap=colors, with_grid=True, with_contour=False, left_string='Vorticity (1/s)', levels=levels, add_cyclic_point=False)
		if f.time.size > 100:
			plot_contour_map(ax, f.vor[100,0,:,:], cmap=colors, with_grid=True, with_contour=False, left_string='Vorticity (1/s)', levels=levels, add_cyclic_point=False)

	ax = fig.add_subplot(2, 1, 2)
	plot_time_series(ax, f.time, abs(f.vor).max(dim=['lev', 'ilat', 'ilon']), ylim=(6.0e-5, 8.0e-5))

	f.close()

if args.loop:
	loop_plot(args.input, fig, plot)
else:
	plot(args.input, fig)
	plt.tight_layout()
	if args.output:
		plt.savefig(args.output, format='pdf')
	else:
		plt.show()
